# Copyright (c) 2018 Presto Labs Pte. Ltd.
# Author: jhkim

import datetime
import zipfile
import os
from absl.flags import DEFINE_string

DEFINE_string('holiday_json_path',
              '/remote/iosg/data/repository/presto_data/recent/presto/data/holidays.json',
              '')

DEFINE_string('corp_action_path',
              '/remote/iosg/data/repository/presto_data/' + 'recent/presto/data/corp_actions.json',
              '')
DEFINE_string(
    'corp_action_path_us',
    '/remote/iosg/data/repository/presto_data/' + 'recent/presto/data/corp_actions_us.json',
    '')


class DataUtil(object):
  # Flags file input status
  FileInputStatus = "file_input_status"
  HolidayCached = None

  def __init__(self):
    pass

  @staticmethod
  def create_updated_table(db_manager):
    table_name = DataUtil.FileInputStatus
    if not db_manager.does_table_exist(table_name):
      query_str = """
CREATE TABLE IF NOT EXISTS %s(
id int auto_increment primary key,
filepath char(255),
status tinyint(1),
insertTime datetime)""" % table_name
      db_manager.execute_raw_query(query_str)

  @staticmethod
  def mark_update_done(db_manager, file_path):
    DataUtil.create_updated_table(db_manager)
    query_str = """
INSERT INTO file_input_status(filepath, status, insertTime)
VALUES('%s', 1, NOW())
""" % (file_path)
    db_manager.execute_raw_query(query_str)

  @staticmethod
  def get_update_done(db_manager, file_path):
    DataUtil.create_updated_table(db_manager)
    selected = db_manager.select_data(DataUtil.FileInputStatus, ['status'], [{
        'filepath': file_path
    }],
                                      like=False,
                                      optional_str="")
    if len(selected) > 0:
      return selected[0]['status']
    else:
      return False

  @staticmethod
  def unzip_file(from_file, to_file, force_unzip, sub_file=None):
    a_zip = zipfile.ZipFile(from_file)
    if sub_file is None:
      if not os.path.exists(to_file) or force_unzip:
        a_zip.extractall(to_file)
    elif not os.path.exists(os.path.join(to_file, sub_file)) or force_unzip:
      a_zip.extract(sub_file, to_file)

  # TODO(jhkim): This is a same code as factset file manager.
  # remove factset file manager code or something.
  @staticmethod
  def get_file_unzipped(zip_dir, zip_file, unzip_dir, force_unzip, sub_file=None):
    out_path = os.path.join(unzip_dir, zip_file)
    DataUtil.unzip_file(os.path.join(zip_dir, zip_file), out_path, force_unzip, sub_file)
    return out_path

  @staticmethod
  def get_updated_files_like(db_manager, filepath_match_str):
    data_sel = db_manager.select_data(DataUtil.FileInputStatus, ['filepath'], [{
        'filepath': filepath_match_str
    }],
                                      like=True,
                                      optional_str="")
    return [data_row['filepath'] for data_row in data_sel]

  @staticmethod
  def delete_file_status_like(db_manager, filepath_match_str):
    db_manager.delete_data(DataUtil.FileInputStatus, [{'filepath': filepath_match_str}], like=True)

  @staticmethod
  def get_general_map_from_db(db_manager, table_name, map_from, map_to):
    db_manager.set_fetch_how(0)
    data_selected = db_manager.select_data(table_name, [map_from, map_to], [],
                                           like=False,
                                           optional_str="")
    db_manager.set_fetch_how(1)
    seen = set()
    keys = [row[0] for row in data_selected if row[0] is not None]
    for a_key in keys:
      if a_key not in seen:
        seen.add(a_key)
      else:
        raise ValueError("Duplicated Key: %s", a_key)
    return {row[0]: row[1] for row in data_selected}

  @staticmethod
  def get_next_bd(a_date, exchange, forward=True):
    holidays_set = set(DataUtil.get_holidays(exchange))
    dt_iter = datetime.datetime.strptime(str(a_date), "%Y%m%d")
    while True:
      if forward:
        dt_iter += datetime.timedelta(days=1)
      else:
        dt_iter -= datetime.timedelta(days=1)
      dt_int = int(dt_iter.strftime("%Y%m%d"))
      if (dt_iter.weekday() <= 4) and (dt_int not in holidays_set):
        return dt_int

  @staticmethod
  def get_days_next_bd(a_date, exchange, days, forward=True):
    date_iter = a_date
    for _ in xrange(days):
      date_iter = DataUtil.get_next_bd(date_iter, exchange, forward)
    return date_iter

  @staticmethod
  def get_range_bds(date_ge, date_lt, exchange):
    holidays_set = set(DataUtil.get_holidays(exchange))
    dt_iter = datetime.datetime.strptime(str(date_ge), "%Y%m%d")
    dts = []
    while True:
      dt_int = int(dt_iter.strftime("%Y%m%d"))
      if (dt_iter.weekday() <= 4) and (dt_int not in holidays_set):
        if dt_int < date_lt:
          dts.append(dt_int)
        else:
          break
      dt_iter += datetime.timedelta(days=1)
    return dts
